'''
Created on Mar 6, 2010

@author: mkiyer
'''
import logging
import xml.etree.cElementTree as etree

import sampledb3 as sdb

def parse_elem_samplegroup(elem):
    # parse the xml elements
    assert elem.tag == 'samplegroup'
    name_elem = elem.find('name')
    name = name_elem.text if name_elem is not None else None
    color_elem = elem.find('color')
    color = color_elem.text if color_elem is not None else None
    cond_elems = elem.findall('cond')
    conds = []
    if cond_elems is not None:
        for cond_elem in cond_elems:
            cond = {}
            for e in cond_elem.getchildren():
                cond[e.tag] = e.text
            conds.append(cond)
    samples = [x.text for x in elem.findall('sample')]
    #library_ids = [x.text for x in elem.findall('library')]
    return name, color, samples, conds

def parse_samplegroups_xml(xmlfile):
    tree = etree.parse(xmlfile)
    root = tree.getroot()
    assert root.tag == 'samplegroups'    
    for e in root.getchildren():
        yield parse_elem_samplegroup(e)

class SamplePool(object):
    def __init__(self, name, library_ids):
        self.name = name
        self.library_ids = library_ids

class SampleGroup(object):
    def __init__(self, name, color):
        self.name = name
        self.color = color
        self.samples = []

def fetch_sample_pool(sample_name, sdbi, qc_filter):
    sample_library_ids = []
    libraries = sdbi.get_libraries_by_sample_name(sample_name, best=True)
    for library in libraries:
        if qc_filter and (library.qc_status == False):
            logging.debug('library %s: marked QC-FAIL excluded from query' % library.id)
            continue
        sample_library_ids.append(library.id)
    if len(sample_library_ids) > 0:
        return SamplePool(sample_name, sample_library_ids)
    return None

def fetch_sample_groups(sample_xml_file, sdbi, qc_filter):
    sample_groups = []
    # parse XML file and make sample groups
    for sample_group_info in parse_samplegroups_xml(sample_xml_file):
        group_name, group_color, names, conds = sample_group_info        
        sgroup = SampleGroup(group_name, group_color)
        # get the libraries associated with each sample
        matching_sample_names = sdbi.get_sample_names(names, conds)
        for sample_name in matching_sample_names:
            sp = fetch_sample_pool(sample_name, sdbi, qc_filter)
            # add pool to group
            if sp != None:
                sgroup.samples.append(sp)
        # add group to list
        if len(sgroup.samples) > 0:
            sample_groups.append(sgroup)
    return sample_groups

def fetch_sample_properties(sample_groups, sdbi):
    # get sample properties - we do this after building the sample pools 
    # because we only want the properties that apply to the collective 
    # set of samples being analyzed and disregard irrelevant properties
    sample_names = []
    for sgroup in sample_groups:        
        sample_names.extend([sp.name for sp in sgroup.samples])
    propdict = sdbi.get_sample_properties(sample_names)
    # add the properties to each sample
    for sgroup in sample_groups:
        for sp in sgroup.samples:
            sp.props = propdict[sp.name]
    return propdict

def write_sample_pools(outfileobj, sample_xml_file, sample_db_file, qc_filter=True):
    # load sample db
    sdbi = sdb.SampleDBInstance.load(sample_db_file)    
    sample_groups = fetch_sample_groups(sample_xml_file, sdbi, qc_filter)
    sample_props = fetch_sample_properties(sample_groups, sdbi)
    # fix the order of the sample properties
    prop_keys = set()
    for sg in sample_groups:
        for sp in sg.samples:
            prop_keys.update(sample_props[sp.name].keys())
    prop_keys = sorted(prop_keys)
    # choose an order for the rows
    row_header = ['name', 'group', 'color', 'library_ids']
    row_header.extend(prop_keys)
    # build columns of sample pool information 
    cols = []
    for sgroup in sample_groups:
        for sp in sgroup.samples:
            col = [sp.name, sgroup.name, sgroup.color, ','.join(sp.library_ids)]
            for prop in prop_keys:
                col.append(str(sample_props[sp.name][prop]))
            cols.append(col)
    # output a tab-delimited table with samples in columns and
    # sample information in rows
    for i, row in enumerate(row_header):
        s = [row_header[i]]        
        s.extend([c[i] for c in cols])
        outfileobj.write('\t'.join(s) + '\n')

def read_sample_pools(line_iter):
    # first line is sample names
    sample_names = line_iter.next().strip().split('\t')[1:]
    # next is sample groups
    sample_groups = line_iter.next().strip().split('\t')[1:]
    # next is color
    colors = line_iter.next().strip().split('\t')[1:]
    # next is library_ids
    library_ids = [s.split(',') for s in line_iter.next().strip().split('\t')[1:]]
    # now build property list and dictionary
    propdict = {}
    for row in line_iter:
        fields = row.strip().split('\t')
        prop = fields[0]
        prop_values = fields[1:]
        propdict[prop] = prop_values
    sample_pools = type('SamplePoolInfo', (object,), 
                        dict(names=sample_names,
                             groups=sample_groups,
                             colors=colors,
                             library_ids=library_ids,
                             properties=propdict))
    return sample_pools

import unittest
class TestSamplePools(unittest.TestCase):
    
    def testReadWrite(self):
        import StringIO
        xmlstr = '''<samplegroups>
  <samplegroup>
    <name>benign_cl</name>
    <color>purple</color>
    <sample>MCF7</sample>
    <cond>
      <app_type>Transcriptome</app_type>
      <tissue_type>Prostate</tissue_type>
      <diagnosis>Benign Cell Line</diagnosis>
      <treated>FALSE</treated>
    </cond>
  </samplegroup>
</samplegroups>
'''
        correct = '''name	MCF7	PWR-1E	PrEC	PrSMC	RWPE	WPMY-1
group	benign_cl	benign_cl	benign_cl	benign_cl	benign_cl	benign_cl
color	purple	purple	purple	purple	purple	purple
library_ids	mctp_20F65AAXX_8,mctp_20E46AAXX_8,mctp_30LEJAAXX_7	mctp_42B08AAXX_2	mctp_30CYNAAXX_5,mctp_209ENAAXX_8,mctp_314T1AAXX_1,mctp_30351AAXX_7	mctp_314T1AAXX_2	mctp_20E6CAAXX_6,mctp_20E6CAAXX_7,mctp_30L0CAAXX_1,mctp_20E6CAAXX_8,mctp_20F0GAAXX_7,mctp_20F0BAAXX_7,mctp_20F0BAAXX_8,mctp_20F0BAAXX_6,mctp_20F0GAAXX_8,mctp_20F0GAAXX_6	mctp_42B08AAXX_3
cell_line-tissue_origin	Pleural Effusion	None	None	None	None	None
cohort	None	Prostate	None	None	None	Prostate
diagnosis	Localized Cell Line	Benign Cell Line	Benign Cell Line	Benign Cell Line	Benign Cell Line	Benign Cell Line
diagnosis_type	Ductal Carcinoma	None	None	None	None	None
er	Positive	None	None	None	None	None
ethnicity	White	None	None	None	None	None
ets	None	ETS-	None	None	None	None
ets_status	None	NA	None	None	None	None
luminal.basal	Luminal	None	None	None	None	None
matched	FALSE	FALSE	FALSE	FALSE	FALSE	FALSE
oncoseq_sample_name	MCF7	PWR-1E	PrEC	PrSMC	RWPE	WPMY-1
transplant	None	Not Transplant	None	None	None	Not Transplant
treated	FALSE	FALSE	FALSE	FALSE	FALSE	FALSE
tumor_stage	Lu	None	None	None	None	None'''
        input = StringIO.StringIO()
        input.write(xmlstr + '\n')
        input.seek(0)
        sample_db_file = 'sampledb_2010_Mar_06.pickle'
        qc_filter=True
        output = StringIO.StringIO()
        write_sample_pools(output, iter(input), sample_db_file, qc_filter)
        output.seek(0)
        result = read_sample_pools(output)
        correctoutput = StringIO.StringIO(correct)
        correctoutput.seek(0)
        correctresult = read_sample_pools(correctoutput)
        self.assertEqual(result.names, correctresult.names)
        self.assertEqual(result.groups, correctresult.groups)
        self.assertEqual(result.library_ids, correctresult.library_ids)
        self.assertEqual(result.properties, correctresult.properties)

if __name__ == '__main__':
    unittest.main()
